package edu.harvard.s3xyback.edu.harvard.s3xyback.examples.crypto;

import java.io.File;
import java.io.FileInputStream;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.Signature;
import java.security.spec.X509EncodedKeySpec;
import java.util.Date;
import java.util.Random;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;

/**
 * CookieDecoderAESwRSA is used to decrypt AMS cookies that were encrypted with 
 * the AES and RSA algorithms. Only one instance of this class should be necessary
 * under normal circumstance.
 * 
 * @author jkennedy
 *
 */
public class CookieDecoderAESwRSA {

	// key sizes to generate
	private static final int RSA_KEYSIZE = 2048; //bits
	private static final int AES_KEYSIZE = 256; // bits
	private static final int SIG_SIZE = RSA_KEYSIZE/8; // bytes
	
	// algorithms to use
	private static final String RSA_STRING = "RSA/ECB/PKCS1Padding";
	private static final String AES_STRING = "AES/CBC/PKCS5Padding";
	private static final String SIG_STRING = "SHA1withRSA";	
	
	// Keys
	private PublicKey rsaPublicKey; // Public key used by AMS server
	private SecretKey aesSessionKey; // Secret key generated for this cookie
	
	// ciphers
	private Signature rsaSig;
	private Cipher rsaCipher;	
	private Cipher aesCipher;
	
	
	/**
	 * Instantiate the decoder for use with a given public key. This key
	 * will be used for both decryption and verification. 
	 * 
	 * @param key rsa public key bytes
	 * @throws CryptoException
	 */
	public CookieDecoderAESwRSA(byte[] key) throws CryptoException {
		initKey(key);
	}
	
	
	/**
	 * Instantiate the decoder for use with a given public key. This key
	 * will be used for both decryption and verification. 
	 * 
	 * @param keyFile file that contains the public key bytes
	 * @throws CryptoException
	 */
	public CookieDecoderAESwRSA(File keyFile) throws CryptoException { 
		try {
			
			byte[] key = new byte[(int)keyFile.length()];
			new FileInputStream(keyFile).read(key);
			initKey(key);
			
		} catch (Exception e) {
			throw new CryptoException(e);
		}
	}
	
	/**
	 * Parse the RSA key bytes and init the ciphers being used.
	 * 
	 * @param key
	 * @throws CryptoException
	 */
	private void initKey(byte[] key) throws CryptoException {
		try {
			
			// load public key
			X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(key);
			KeyFactory kf = KeyFactory.getInstance("RSA");
			rsaPublicKey = kf.generatePublic(publicKeySpec);			
		
			
			// load and init rsa cipher
			rsaCipher = Cipher.getInstance(RSA_STRING);
			rsaCipher.init(Cipher.UNWRAP_MODE, rsaPublicKey);
			
			// load and init rsa signature
			rsaSig = Signature.getInstance(SIG_STRING);
			rsaSig.initVerify(rsaPublicKey);
			
			// load aes algorithm but can't init until initSession()
			aesCipher = Cipher.getInstance(AES_STRING);
			
		} catch(Exception e) {
			throw new CryptoException(e);
		}
	}
	
	
	/**
	 * Set the encrypted, base64 encoded session key. Will be used for
	 * future calls to decrypt.
	 * 
	 * @param key The key should be encrypted with RSA and base65 encoded as a String
	 */
	public void initSession(String key) throws CryptoException {
		initSession(Base64.decodeToBytes(key));		
	}

	
	/**
	 * Set the encrypted session key. Will be used for
	 * future calls to decrypt.
	 * 
	 * @param key The key should be encrypted with RSA
	 */
	public void initSession(byte[] key) throws CryptoException {

		try {
			
			// decrypt the session key							
			aesSessionKey = (SecretKey)rsaCipher.unwrap(key, "AES", Cipher.SECRET_KEY);
			
			// use session key for IV since it's random
			IvParameterSpec iv = new IvParameterSpec(aesSessionKey.getEncoded(), 0, 16);
			
			// initialize the cipher object			
			aesCipher.init(Cipher.DECRYPT_MODE, aesSessionKey, iv);
			
		} catch (Exception e) {
			throw new CryptoException(e);
		}
	}	
	
	
	/**
	 * Decrypt the message and verify the signature
	 * 
	 * @param msg the encrypted, base64 encoded cookie
	 * @return
	 */
	public String decrypt(String msg) throws CryptoException {
		// decode the base 64		
		return decrypt(Base64.decodeToBytes(msg));
	}
	
	
	/**
	 * Decrypt the message and verify the signature
	 * 
	 * @param msg the encrypted cookie
	 * @return
	 */
	public String decrypt(byte[] encMsg) throws CryptoException {
		try {
			// decrypt the message									
			byte[] msgBytes = aesCipher.doFinal(encMsg);
			
			// parse plaintext and signature
			byte[] sigBytes = subset(msgBytes, 0, SIG_SIZE);
			byte[] plainBytes = subset(msgBytes, SIG_SIZE, msgBytes.length);
			
			// verify the plaintext
			rsaSig.update(plainBytes);
			if (! rsaSig.verify(sigBytes))
				throw new CryptoException("Signature was invalid");
	
			// return the decrypted message without signature
			return new String(plainBytes, "UTF8");
			
		} catch (Exception e) {
			throw new CryptoException(e);
		}
		
	}	
	
	
	/**
	 * Return a subset of the byte array as a new byte array
	 * @param bytes
	 * @param start
	 * @param end
	 * @return
	 */
	private static byte[] subset(byte[] bytes, int start, int end) {
		int length = end-start;
		if (length < 1) throw new IllegalArgumentException("Can not end before start");
		byte[] buff = new byte[length];
		for (int i = 0; i < length; i++) {
			buff[i] = bytes[i+start];
		}
		return buff;
	}
	
	/**
	 * Generate a random string, used for main() performance test
	 * @param length
	 * @return
	 */
	private static String getRandomString(int length) {
		String chars = "abcdefghijklmonpqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789";
		Random rand = new Random();
		char[] buf = new char[length];
		for (int i = 0; i < buf.length; i++) {
			buf[i] = chars.charAt(rand.nextInt(chars.length()));
		}
		return new String(buf);
	}	
	
	
	/**
	 * 
	 */
	public static void main(String args[]) throws Exception {

		KeyPairGenerator rsaKeyGen = KeyPairGenerator.getInstance("RSA");		
		rsaKeyGen.initialize(RSA_KEYSIZE);	
		KeyPair pair = rsaKeyGen.generateKeyPair();
		PublicKey rsaPub = pair.getPublic();
		PrivateKey rsaPriv = pair.getPrivate();		
		
		CookieEncoderAESwRSA encoder = new CookieEncoderAESwRSA(rsaPriv.getEncoded());
		CookieDecoderAESwRSA decoder = new CookieDecoderAESwRSA(rsaPub.getEncoded());
		
		// vars for tracking timing metrics
		Date startDt = null;
		Date endDt = null;
		long time = 0;
		long encTime = 0;
		long encHigh = 0;
		long decTime = 0;
		long decHigh = 0;		
		
		// run the tests
		int i;		
		for (i = 0; i < 100; i++) {	
			String plaintext = getRandomString(300);
			String plaintext2 = null;
			String ciphertext = null;
			
			//encrypt			
			startDt = new Date();
			String skey = encoder.initSession();
			ciphertext = encoder.encrypt(plaintext);
			endDt = new Date();
			
			// calc time metrics
			time = endDt.getTime() - startDt.getTime();
			encTime += time;
			if (time > encHigh) encHigh = time;			
			
			// Decrypt the ciphertext
			startDt = new Date();
			decoder.initSession(skey);
			plaintext2 = decoder.decrypt(ciphertext);
			endDt = new Date();
			
			// calc time metrics
			time = endDt.getTime() - startDt.getTime();
			decTime += time;
			if (time > decHigh) decHigh = time;
			
			if (! plaintext.equals(plaintext2))
				throw new Exception("Plaintext mismatch: " + plaintext + " != " + plaintext2);
		}
		
		// Drop the highest time to account for initial class loading
		System.out.println("Dropping encryption high time: " + encHigh);
		System.out.println("Average encryption time: " + ((encTime - encHigh)/i));
		System.out.println("Dropping decryption high time: " + decHigh);
		System.out.println("Average decryption time: " + ((decTime - decHigh)/i));		
	}
}
